import time
import win32gui, win32ui, win32con, win32api
import numpy as np
import cv2
from PIL import ImageGrab

screen_width,screen_height=1920,1080
grab_width,grab_height = 1920,1080


def my_print(_obj,_obj_name=''):
    print(f'\n{_obj_name}:\n{_obj}\n')

hwnd_title = dict()
def get_all_hwnd(hwnd,mouse):
    if win32gui.IsWindow(hwnd) and win32gui.IsWindowEnabled(hwnd) and win32gui.IsWindowVisible(hwnd):
        hwnd_title.update({hwnd:win32gui.GetWindowText(hwnd)})

def update_hwnd_title():
    win32gui.EnumWindows(get_all_hwnd,0)
    for h,t in hwnd_title.items():
        if t is not "":
            print(h,t)


def grab_screen_win32_v2(window_title,grab_rect=None):
    hwnd = 0  # 窗口的编号，0号表示当前活跃窗口
    hwnd = win32gui.FindWindow(None,window_title)
    if hwnd==0:
        print(f'{window_title} 窗口可能未找到')
    # 根据窗口句柄获取窗口的设备上下文DC（Divice Context）
    hwndDC = win32gui.GetWindowDC(hwnd)
    # 根据窗口的DC获取mfcDC
    mfcDC = win32ui.CreateDCFromHandle(hwndDC)
    # mfcDC创建可兼容的DC
    saveDC = mfcDC.CreateCompatibleDC()
    # 创建bigmap准备保存图片
    saveBitMap = win32ui.CreateBitmap()
    x, y, w, h = (0,0,100,100)
    if grab_rect is None:
        # 获取监控器信息
        MoniterDev = win32api.EnumDisplayMonitors(None, None)
        w = MoniterDev[0][2][2]
        h = MoniterDev[0][2][3]
        # print w,h　　　#图片大小
        # 为bitmap开辟空间
        saveBitMap.CreateCompatibleBitmap(mfcDC, w, h)
        # 高度saveDC，将截图保存到saveBitmap中
        saveDC.SelectObject(saveBitMap)
    else:
        x,y,w,h = grab_rect
        saveBitMap.CreateCompatibleBitmap(mfcDC,w,h)
        saveDC.SelectObject(saveBitMap)

    # 截取从左上角（0，0）长宽为（w，h）的图片
    saveDC.BitBlt((0, 0), (w, h), mfcDC, (x, y), win32con.SRCCOPY)
    signed_ints_array = saveBitMap.GetBitmapBits(True)
    img = np.frombuffer(signed_ints_array,dtype="uint8")
    img.shape = (h,w,4)
    img = cv2.cvtColor(img,cv2.COLOR_RGBA2RGB)

    win32gui.DeleteObject(saveBitMap.GetHandle())
    mfcDC.DeleteDC()
    saveDC.DeleteDC()

    return img

def main():
    time_tag=time.time()
    count=0;
    id=0;
    while True:
        count+=1

        img = grab_screen_win32_v2(window_title='Counter-Strike: Global Offensive - Direct3D 9',grab_rect=(0,0,grab_width,grab_height))
        if count>=500:
            count=0
            id+=1
            cv2.imwrite(f'D:\PYPRJ\yolov5-7.0\data\mydata2\{id}.jpg',img)
        cv2.imshow('grab_win32',img)
        cv2.waitKey(1)

    print(f'用时：{time.time()-time_tag}  FPS:{1/(time.time()-time_tag)}')

if __name__ == '__main__':
    main()